import torch
from torch import nn
import torch.nn.functional as F



class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, 
            kernel_size, stride=1, padding=0, 
            bias=True,
            norm_func=None
    ):
        super().__init__()

        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=bias
        )

        self.norm = norm_func

        for m in [self.conv, self.norm]:
            if isinstance(m, nn.Conv2d):
                torch.nn.init.normal_(m.weight, mean=0, std=0.01)
                torch.nn.init.constant_(m.bias, 0)
            if isinstance(m, nn.GroupNorm):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.conv(x)
        
        if self.norm is not None:
            x = self.norm(x)
        
        return F.relu(x)

